import asyncio
from .通用 import UDP服务器信息, 读取参数
from .UDP客户端 import open_connection
from .UDP消息 import UDP消息


async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
    udp客户端 = await open_connection(UDP服务器信息.服务器地址, UDP服务器信息.服务器端口)
    if udp客户端 is None:
        return

    async def 读取任务():
        while not writer.is_closing() and not udp客户端.会话.会话关闭:
            try:
                data = await reader.read(1200)
            except:
                writer.close()
                udp客户端.关闭会话()
                return
            if not data:
                writer.close()
                udp客户端.关闭会话()
            try:
                await udp客户端.会话.发送数据(data)
            except:
                return

    async def 写回任务():
        while not writer.is_closing() and not udp客户端.会话.会话关闭:
            try:
                消息: UDP消息 = await udp客户端.会话.读一条消息()
                writer.write(消息.内容)
            except:
                writer.close()
                return

    await asyncio.gather(读取任务(), 写回任务())


async def 启动隧道本地TCP服务器():
    server = await asyncio.start_server(handler, UDP服务器信息.监听地址, UDP服务器信息.监听端口)
    try:
        async with server:
            await server.serve_forever()
    except Exception:
        pass


def main():
    读取参数()
    asyncio.run(启动隧道本地TCP服务器())